# -*- coding: utf-8 -*-

'''
@project : SmartHomeVideo
@FileName: video_analysis_main
@Author  :linych
@Time    :2021/2/22 18:30
@Desc  :
'''

import os
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import pika
import json
from VideoAnalysis.Download import video_load
from Core.redis_connecter import redis_Task, redis_Param, TASK_NAME
from Core import conf
import cv2
import base64
from MQ.mq_connect_util import MQ_connect

host = conf.get('MQ', 'url')
username = conf.get('MQ', 'username')
pwd = conf.get('MQ', 'pwd')
queue = conf.get('MQ', 'queue')


import asyncio

from io import BytesIO
import numpy as  np
from multiprocessing import Queue, Process

VIDEO_INPUT_LIST = Queue()

loop = asyncio.new_event_loop()

# 建立连接，服务器地址为localhost，可指定ip地址
print(host)


async def video_download(url, correlation_id, routing_key):
    s = time.time()
    out = await video_load.img_download(url)
    data = (out, correlation_id, routing_key)
    VIDEO_INPUT_LIST.put(data)

    print('video_download use : {}'.format(time.time() - s))
    return out


def video_decoder(vinput):
    s = time.time()

    if vinput is None:
        return []

    vinput = BytesIO(vinput)

    from av import open
    video = open(vinput)

    stream = next(s for s in video.streams if s.type == 'video')
    stream.codec_context.skip_frame = 'NONKEY'
    video_list = []
    for packet in video.demux(stream):
        for frame in packet.decode():
            if frame.key_frame:
                img = frame.to_ndarray(format='rgb24')
                video_list.append(img)

    res_list = []

    for i in range(len(video_list)):
        retval, buffer = cv2.imencode('.jpg', video_list[i])
        pic_str = base64.b64encode(buffer)
        pic_str = pic_str.decode()

        res_list.append(pic_str)

    print('video_decoder {} use: {}'.format(np.array(video_list).shape, time.time() - s))

    return res_list


def result_to_redis(correlation_id, routing_key, result):
    try:

        s = time.time()

        redis_Param.hmset(correlation_id, {'routing_key': routing_key, 'correlation_id': correlation_id})
        redis_Param.expire(correlation_id, 5 * 60)

        # redis存图片
        for i in range(len(result)):
            tm_key = correlation_id + '_{}'.format(i)

            # redis记录抽帧图片base64
            redis_Task.set(tm_key, result[i])
            redis_Task.expire(tm_key, 5 * 60)

            # 记录映射
            redis_Task.rpush(correlation_id, tm_key)

            redis_Task.expire(correlation_id, 5 * 60)

        redis_Task.rpush(TASK_NAME, correlation_id)

        print('result_to_redis {} use {}'.format(correlation_id, time.time() - s))

    except Exception as e:

        print('result_to_redis error {}'.format(e))


def decoder_task(que):
    Decoder_Time_list = []

    while True:

        try:

            try:
                print('Queue size : {}'.format(que.qsize()))
                data = que.get()
            except Exception as e:
                print(e)
                data = None

            if data is None:
                # print('wait..........................................')
                time.sleep(0.001)
                continue

            s = time.time()

            vinput, correlation_id, routing_key = data

            result = video_decoder(vinput=vinput)

            t = threading.Thread(target=result_to_redis, args=(correlation_id, routing_key, result))
            t.start()

            PID = os.getpid()
            Decoder_Time_list.append(time.time() - s)

            print('decoder_task  {}...avg : {}  len:{}'.format(PID, sum(Decoder_Time_list) / len(Decoder_Time_list),
                                                               len(Decoder_Time_list)))

        except Exception as e:

            print('decoder_task error : {}'.format(e))


# 对RPC请求队列中的请求进行处理
def on_request(ch, method, props, body):
    global VIDEO_INPUT_LIST

    # print('consume doing.....')
    arg_dict = json.loads(body)
    url = arg_dict.get('fileurl', '')

    routing_key = props.reply_to
    correlation_id = props.correlation_id

    while VIDEO_INPUT_LIST.qsize() > 100:
        print('waiting...{}'.format(VIDEO_INPUT_LIST.qsize()))
        time.sleep(0.1)

    try:

        asyncio.run_coroutine_threadsafe(video_download(url, correlation_id, routing_key), loop)


    except Exception as e:
        print(e)

    ch.basic_ack(delivery_tag=method.delivery_tag)


def video_main():

    while True:
        try:
            print('begin.............')

            mq_client = MQ_connect(host=host, username=username, pwd=pwd, queue=queue)

            channel = mq_client.get_channel()

            # 负载均衡，同一时刻发送给该服务器的请求不超过一个
            # channel.basic_qos(prefetch_count=10)

            channel.basic_consume(on_message_callback=on_request, queue=queue)

            channel.start_consuming()

        except Exception as e:

            print('consume error : {}'.format(e))


def start_loop(loop):
    asyncio.set_event_loop(loop)
    print("start loop", time.time())
    loop.run_forever()
    print(asyncio.Task.all_tasks().pop().result())


if __name__ == '__main__':

    import threading

    download_worker = int(conf.get('VIDEO','download_worker'))
    decoder_worker = int(conf.get('VIDEO','decoder_worker'))



    threads = []


    for i in range(download_worker):
        t = threading.Thread(target=video_main)
        threads.append(t)
        t.start()

    # 解码
    l = []

    for i in range(decoder_worker):
        t = Process(target=decoder_task, args=(VIDEO_INPUT_LIST,))
        l.append(t)
        t.start()

    start_loop(loop=loop)
